#!/usr/bin/env python
"""
微信聊天助手冒烟测试脚本
使用真实的微信终端和数据库进行基本功能验证
"""
import os
import sys
import time
import json
import logging
import asyncio
import datetime
import traceback
from typing import List, Dict, Any, Optional

# 确保能够导入应用模块
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

# 设置环境变量
from tests.env_setup import setup_test_environment
setup_test_environment()

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select
from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.asyncio import create_async_engine

# 导入测试配置
from tests.smoke_test_config import SMOKE_TEST_CONFIG, TEST_USERS, PUA_TEST_PATTERNS

# 应用程序导入
from app.core.config import settings
from app.db.session import get_db
from app.models.contacts import Contact
from app.models.messages import Message, MessageDirection, MessageType, MessageStatus
from app.models.sessions import Session
from app.models.pua import PUAPattern, PUAPatternCategory
from app.services.wechat import wxauto_service
from app.services.pua_detection import pua_detector

# 配置日志
log_config = SMOKE_TEST_CONFIG["logging"]
log_handlers = []
if log_config.get("console", True):
    log_handlers.append(logging.StreamHandler())
if log_config.get("file"):
    os.makedirs(os.path.dirname(log_config["file"]), exist_ok=True)
    log_handlers.append(logging.FileHandler(log_config["file"]))

logging.basicConfig(
    level=log_config.get("level", logging.INFO),
    format=log_config.get("format", "%(asctime)s - %(name)s - %(levelname)s - %(message)s"),
    handlers=log_handlers
)
logger = logging.getLogger("smoke_test")

# 创建真实数据库连接（非内存数据库）
DATABASE_URL = f"sqlite+aiosqlite:///{SMOKE_TEST_CONFIG['database']['url']}"

# 创建异步引擎
engine = create_async_engine(
    DATABASE_URL,
    echo=SMOKE_TEST_CONFIG["database"]["echo"],
    # 添加SQLite特定配置
    connect_args={"check_same_thread": False}
)

# 创建异步会话工厂
async_session = sessionmaker(
    engine,
    class_=AsyncSession,
    expire_on_commit=False,
    autocommit=False,
    autoflush=False
)

class SmokeTest:
    """微信聊天助手冒烟测试类"""

    def __init__(self):
        """初始化测试环境"""
        self.session = None
        self.wxauto = None
        self.config = SMOKE_TEST_CONFIG

    async def setup(self):
        """设置测试环境"""
        try:
            # 创建数据库会话
            self.session = async_session()
            logger.info("数据库连接已初始化")

            # 初始化微信服务
            self.wxauto = wxauto_service
            logger.info("微信服务已初始化")
            
            # 确保数据库表存在
            if self.config["database"]["create_tables"]:
                from app.db.base import Base
                async with engine.begin() as conn:
                    await conn.run_sync(Base.metadata.create_all)
                logger.info("数据库表已创建")

            return True
        except Exception as e:
            logger.error(f"设置测试环境失败: {str(e)}")
            return False

    async def test_wechat_connection(self) -> bool:
        """测试微信连接"""
        try:
            # 首先确认服务实例可用
            if not self.wxauto:
                logger.error("微信服务实例不可用")
                return False
            
            # 检查微信是否已登录
            if self.wxauto.check_login():
                logger.info("微信已登录")
                return True
            else:
                logger.error("微信未登录，请先登录微信")
                return False
        except Exception as e:
            logger.error(f"微信连接测试失败: {e}")
            logger.error(f"异常详情: {traceback.format_exc()}")
            return False

    async def test_database_connection(self) -> bool:
        """测试数据库连接"""
        try:
            # 执行简单查询
            result = self.session.execute(text("SELECT 1"))
            if result:
                logger.info("数据库连接正常")
                return True
            else:
                logger.error("数据库连接测试失败")
                return False
        except Exception as e:
            logger.error(f"数据库连接测试异常: {e}")
            return False

    async def test_chat_sessions(self) -> bool:
        """测试获取微信聊天会话列表"""
        try:
            sessions = self.wxauto.get_session_list()
            if sessions and len(sessions) > 0:
                logger.info(f"成功获取聊天会话，共 {len(sessions)} 个")
                logger.info(f"部分会话名称: {sessions[:3]}")
                return True
            else:
                logger.warning("未获取到聊天会话")
                return False
        except Exception as e:
            logger.error(f"获取聊天会话失败: {e}")
            return False

    async def test_messages_retrieval(self, session_name: Optional[str] = None) -> bool:
        """测试获取微信消息"""
        try:
            # 如果未指定会话，则获取第一个会话或使用配置中的会话
            if not session_name:
                session_name = self.config["testing"].get("test_session_name")
                if not session_name:
                    sessions = self.wxauto.get_session_list()
                    if not sessions:
                        logger.error("未找到任何会话")
                        return False
                    session_name = sessions[0]
                
            logger.info(f"尝试获取会话 '{session_name}' 的消息")
            
            # 切换到指定会话
            if self.wxauto.chat_with(session_name):
                # 获取消息
                messages = self.wxauto.get_message_list()
                if messages and len(messages) > 0:
                    logger.info(f"成功获取消息，共 {len(messages)} 条")
                    # 显示最近的一条消息内容
                    if len(messages) > 0:
                        latest_msg = messages[-1]
                        content = latest_msg.get('content', '无内容')
                        logger.info(f"最新消息: {content[:50]}..." if len(content) > 50 else content)
                    return True
                else:
                    logger.warning(f"会话 '{session_name}' 中未获取到消息")
                    return False
            else:
                logger.error(f"切换到会话 '{session_name}' 失败")
                return False
        except Exception as e:
            logger.error(f"获取消息失败: {e}")
            return False

    async def setup_test_pua_patterns(self) -> bool:
        """设置测试用的PUA模式数据"""
        try:
            # 检查是否已存在PUA模式
            result = self.session.execute(text("SELECT COUNT(*) FROM puapattern"))
            count = result.scalar()
            
            if count > 0:
                logger.info(f"数据库中已存在 {count} 条PUA模式记录")
                return True
            
            # 添加测试PUA模式
            for pattern_data in PUA_TEST_PATTERNS:
                pattern = PUAPattern(
                    pattern=pattern_data["pattern"],
                    category=pattern_data["category"],
                    description=pattern_data["description"],
                    severity=pattern_data["severity"]
                )
                self.session.add(pattern)
            
            self.session.commit()
            logger.info(f"成功添加 {len(PUA_TEST_PATTERNS)} 条PUA测试模式")
            return True
            
        except Exception as e:
            logger.error(f"设置PUA测试数据失败: {e}")
            self.session.rollback()
            return False

    async def test_pua_detection(self, text: Optional[str] = None) -> bool:
        """测试PUA检测功能"""
        try:
            # 如果未指定测试文本，则使用配置中的第一条文本
            if not text:
                text = self.config["testing"]["pua_test_texts"][0]
                
            # 确保PUA模式数据已设置
            await self.setup_test_pua_patterns()
                
            logger.info(f"测试PUA检测，文本: {text}")
            result = await pua_detector.detect_pua(text)
            
            logger.info(f"PUA检测结果: 是否为PUA={result.get('is_pua', False)}")
            if 'score' in result:
                logger.info(f"PUA评分: {result['score']}")
            if 'detected_patterns' in result and result['detected_patterns']:
                patterns = [p.get('category', 'unknown') for p in result['detected_patterns']]
                logger.info(f"检测到的模式: {patterns}")
            
            return True
        except Exception as e:
            logger.error(f"PUA检测测试失败: {e}")
            return False

    async def test_message_sending(self, session_name: Optional[str] = None, message: Optional[str] = None) -> bool:
        """测试发送消息功能 (谨慎使用)"""
        if not self.config["testing"].get("run_message_sending", False):
            logger.info("消息发送测试已在配置中禁用")
            return True
            
        try:
            # 如果未指定会话，则获取第一个会话或使用配置中的会话
            if not session_name:
                session_name = self.config["testing"].get("test_session_name")
                if not session_name:
                    sessions = self.wxauto.get_session_list()
                    if not sessions:
                        logger.error("未找到任何会话")
                        return False
                    session_name = sessions[0]
            
            # 如果未指定消息，则使用配置中的测试消息
            if not message:
                message = self.config["testing"].get("test_message", "这是一条测试消息，请忽略")
            
            confirm = input(f"确定要向 '{session_name}' 发送测试消息吗? (y/n): ")
            if confirm.lower() != 'y':
                logger.info("已取消发送测试消息")
                return False
                
            logger.info(f"尝试向会话 '{session_name}' 发送消息: {message}")
            
            # 切换到指定会话
            if self.wxauto.chat_with(session_name):
                # 发送消息
                success = self.wxauto.send_message(message)
                if success:
                    logger.info("消息发送成功")
                    return True
                else:
                    logger.error("消息发送失败")
                    return False
            else:
                logger.error(f"切换到会话 '{session_name}' 失败")
                return False
        except Exception as e:
            logger.error(f"发送消息失败: {e}")
            return False

    async def test_data_saving(self, session_name: Optional[str] = None) -> bool:
        """测试将消息保存到数据库"""
        try:
            # 如果未指定会话，则获取第一个会话或使用配置中的会话
            if not session_name:
                session_name = self.config["testing"].get("test_session_name")
                if not session_name:
                    sessions = self.wxauto.get_session_list()
                    if not sessions:
                        logger.error("未找到任何会话")
                        return False
                    session_name = sessions[0]
                    logger.info(f"自动选择第一个会话: {session_name}")
            
            logger.info(f"尝试保存会话 '{session_name}' 的消息到数据库")
            
            # 切换到指定会话
            logger.info(f"正在切换到会话: {session_name}")
            if not self.wxauto.chat_with(session_name):
                logger.error(f"切换到会话 '{session_name}' 失败")
                return False
            
            # 获取消息
            logger.info("正在获取消息列表...")
            messages = self.wxauto.get_message_list()
            
            # 即使没有消息，也创建会话和联系人记录
            if not messages:
                logger.warning(f"会话 '{session_name}' 中未获取到消息，将只创建会话和联系人记录")
            else:
                logger.info(f"获取到 {len(messages)} 条消息，准备保存到数据库")
                for i, msg in enumerate(messages[:3]):  # 只打印前3条消息作为样例
                    logger.debug(f"消息样例 {i+1}: {msg}")
            
            try:
                # 创建会话记录
                logger.info("正在查找或创建会话记录...")
                
                # 先检查是否已存在该会话
                stmt = select(Session).where(Session.name == session_name)
                result = await self.session.execute(stmt)
                existing_session = result.scalars().first()
                
                is_new_session = False
                if existing_session:
                    logger.info(f"会话 '{session_name}' 已存在，ID: {existing_session.id}")
                    session_obj = existing_session
                else:
                    # 创建新会话
                    is_new_session = True
                    session_obj = Session(
                        name=session_name,
                        session_type="individual" if not session_name.endswith("群聊") else "group"
                    )
                    self.session.add(session_obj)
                    await self.session.flush()  # 刷新以获取ID
                    logger.info(f"会话记录创建成功，ID: {session_obj.id}")
                
                # 创建联系人记录
                logger.info("正在查找或创建联系人记录...")
                contact_name = session_name
                
                # 先检查是否已存在该联系人
                stmt = select(Contact).where(Contact.wechat_id == contact_name)
                result = await self.session.execute(stmt)
                existing_contact = result.scalars().first()
                
                if existing_contact:
                    logger.info(f"联系人 '{contact_name}' 已存在，ID: {existing_contact.id}")
                    contact_obj = existing_contact
                else:
                    # 创建新联系人
                    contact_obj = Contact(
                        wechat_id=contact_name,
                        nickname=contact_name,
                        relationship_type="other"
                    )
                    self.session.add(contact_obj)
                    await self.session.flush()  # 刷新以获取ID
                    logger.info(f"联系人记录创建成功，ID: {contact_obj.id}")
                
                # 如果有消息，进行增量保存
                saved_count = 0
                skipped_count = 0
                failed_count = 0
                
                if messages:
                    # 如果会话已存在，先获取数据库中已有的消息
                    existing_messages = []
                    message_keys = set()
                    
                    if not is_new_session:
                        logger.info("检查已有消息，进行增量保存...")
                        # 查询该会话下的所有消息
                        messages_query = select(Message).where(Message.session_id == session_obj.id)
                        existing_messages_result = await self.session.execute(messages_query)
                        existing_messages = existing_messages_result.scalars().all()
                        
                        # 创建已有消息的唯一键集合，用于后续增量保存
                        for msg in existing_messages:
                            # 使用内容和发送时间作为唯一标识
                            msg_content = msg.content[:50] if msg.content else ""  # 取内容前50个字符
                            msg_key = f"{msg.direction}:{msg.sent_at.isoformat()}:{msg_content}"
                            message_keys.add(msg_key)
                        
                        logger.info(f"数据库中已有 {len(existing_messages)} 条消息")
                    
                    # 限制处理的消息数量
                    max_messages = self.config["testing"].get("max_messages_to_save", 10)
                    messages_to_save = messages[:max_messages] if len(messages) > max_messages else messages
                    logger.info(f"将检查 {len(messages_to_save)} 条新消息（最大限制：{max_messages}条）")
                    
                    for i, msg in enumerate(messages_to_save, 1):
                        try:
                            # 确保消息是字典格式
                            if not isinstance(msg, dict):
                                logger.warning(f"跳过非字典格式的消息 {i}: {type(msg)}")
                                continue
                            
                            # 检查消息是否为None
                            if msg is None:
                                logger.warning(f"跳过空消息 {i}")
                                continue
                            
                            # 记录当前正在处理的消息
                            logger.debug(f"正在处理第 {i} 条消息: {msg}")
                            
                            # 检查必要字段
                            required_fields = ["direction", "msg_type", "content", "sent_at", "status"]
                            missing_fields = [field for field in required_fields if field not in msg]
                            if missing_fields:
                                logger.warning(f"消息 {i} 缺少必要字段: {missing_fields}")
                                msg.update({field: None for field in missing_fields})
                            
                            # 检查该消息是否已存在（增量保存）
                            if not is_new_session:
                                content_preview = msg.get("content", "")[:50]
                                sent_at = msg.get("sent_at")
                                if isinstance(sent_at, datetime.datetime):
                                    sent_at_str = sent_at.isoformat()
                                else:
                                    sent_at_str = str(sent_at)
                                
                                msg_key = f"{msg.get('direction')}:{sent_at_str}:{content_preview}"
                                
                                if msg_key in message_keys:
                                    logger.debug(f"跳过已存在的消息: {content_preview[:20]}...")
                                    skipped_count += 1
                                    continue
                            
                            # 创建消息对象
                            msg_obj = Message(
                                contact_id=contact_obj.id,
                                session_id=session_obj.id,
                                direction=msg.get("direction", MessageDirection.INCOMING),
                                msg_type=msg.get("msg_type", MessageType.TEXT),
                                content=msg.get("content", ""),
                                sent_at=msg.get("sent_at", datetime.datetime.now()),
                                status=msg.get("status", MessageStatus.RECEIVED),
                                extra_data=json.dumps(msg.get("extra_data", {})) if msg.get("extra_data") else None
                            )
                            
                            # 处理媒体文件信息
                            if msg.get("extra_data") and msg.get("extra_data").get("media"):
                                media_info = msg.get("extra_data").get("media")
                                msg_obj.media_path = media_info.get("path")
                                msg_obj.media_type = media_info.get("type")
                                msg_obj.media_size = media_info.get("size")
                                msg_obj.media_name = media_info.get("filename") or media_info.get("original_filename")
                            
                            # 记录消息对象的属性
                            logger.debug(f"消息对象属性: direction={msg_obj.direction}, "
                                       f"msg_type={msg_obj.msg_type}, status={msg_obj.status}, "
                                       f"content={msg_obj.content[:50]}...")
                            
                            self.session.add(msg_obj)
                            saved_count += 1
                            
                            # 每10条消息刷新一次
                            if saved_count % 10 == 0:
                                await self.session.flush()
                                logger.info(f"已成功保存 {saved_count} 条消息")
                                
                        except Exception as e:
                            failed_count += 1
                            logger.error(f"保存第 {i} 条消息失败: {str(e)}")
                            logger.error(f"异常详情: {traceback.format_exc()}")
                            continue
                else:
                    logger.info("没有消息需要保存，只保存会话和联系人信息")
                
                # 更新会话的最后活跃时间
                session_obj.last_active = datetime.datetime.now()
                
                # 提交所有更改
                await self.session.commit()
                
                if messages:
                    logger.info(f"消息保存完成: 新增 {saved_count} 条, 跳过 {skipped_count} 条, 失败 {failed_count} 条")
                    return saved_count > 0 or skipped_count > 0
                else:
                    logger.info("会话和联系人信息保存成功")
                    return True  # 即使没有消息，也认为保存成功
                
            except Exception as e:
                logger.error(f"保存消息过程中发生错误: {str(e)}")
                logger.error(f"异常详情: {traceback.format_exc()}")
                await self.session.rollback()
                return False
                
        except Exception as e:
            logger.error(f"保存消息到数据库失败: {str(e)}")
            logger.error(f"异常详情: {traceback.format_exc()}")
            return False

    async def cleanup(self):
        """清理测试环境"""
        try:
            if self.session:
                await self.session.close()
                logger.info("数据库会话已关闭")
        except Exception as e:
            logger.error(f"清理环境失败: {e}")


async def run_smoke_tests():
    """运行所有冒烟测试"""
    logger.info("开始执行微信聊天助手冒烟测试")
    logger.info(f"数据库URL: {DATABASE_URL}")
    
    test = SmokeTest()
    
    # 设置测试环境
    logger.info("正在设置测试环境...")
    setup_success = await test.setup()
    if not setup_success:
        logger.error("测试环境设置失败，中止测试")
        return
    
    # 记录测试结果
    results = {}
    
    try:
        # 测试微信连接
        logger.info("\n===== 测试微信连接 =====")
        results["微信连接"] = await test.test_wechat_connection()
        
        # 测试数据库连接
        logger.info("\n===== 测试数据库连接 =====")
        results["数据库连接"] = await test.test_database_connection()
        
        # 如果基础连接测试通过，继续进行功能测试
        if results["微信连接"] and results["数据库连接"]:
            # 测试获取聊天会话
            logger.info("\n===== 测试获取聊天会话 =====")
            results["获取聊天会话"] = await test.test_chat_sessions()
            
            # 测试获取消息
            logger.info("\n===== 测试获取消息 =====")
            results["获取消息"] = await test.test_messages_retrieval()
            
            # 测试PUA检测功能
            logger.info("\n===== 测试PUA检测 =====")
            results["PUA检测"] = await test.test_pua_detection()
            
            # 测试保存数据到数据库
            logger.info("\n===== 测试数据保存 =====")
            results["数据保存"] = await test.test_data_saving()
            
            # 测试发送消息
            if SMOKE_TEST_CONFIG["testing"].get("run_message_sending", False):
                logger.info("\n===== 测试消息发送 =====")
                results["发送消息"] = await test.test_message_sending()
        
    finally:
        # 无论测试成功或失败，都要清理环境
        logger.info("\n===== 清理测试环境 =====")
        await test.cleanup()
    
    # 打印测试结果摘要
    logger.info("\n===== 测试结果摘要 =====")
    all_passed = True
    for name, passed in results.items():
        status = "通过" if passed else "失败"
        logger.info(f"{name}: {status}")
        if not passed:
            all_passed = False
    
    logger.info(f"\n总体结果: {'所有测试通过' if all_passed else '部分测试失败'}")
    
    return all_passed


if __name__ == "__main__":
    # 运行测试
    asyncio.run(run_smoke_tests()) 